#include <stdio.h>
#include <stdlib.h>
#include <libavcodec/avcodec.h>
#include <libavformat/avformat.h>
#include <libavutil/opt.h>

typedef struct MediaContext {
    AVFormatContext* fmt_ctx;
    AVCodec* v_codec;
    AVCodec* a_codec;
    AVCodecContext* v_codec_ctx;
    AVCodecContext* a_codec_ctx;
    AVStream* v_stm;
    AVStream* a_stm;
    int v_idx;
    int a_idx;
    char* name;
} MediaContext;

typedef struct MediaDict{
    char copy_video;
    char copy_audio;
    char *codec_opt_key;
    char *codec_opt_value;
    char *muxer_opt_key;
    char *muxer_opt_value;
    char *audio_codec;
    char *video_codec;
    char *extension;  
} MediaDict;

// 为解码器AVFormatContext寻找信息
static int open_media(AVFormatContext **fmt_ctx, const char *filename){
    int ret = 0;
    *fmt_ctx = avformat_alloc_context();
    if((*fmt_ctx) == NULL){
        av_log(NULL, AV_LOG_ERROR, "avformat_alloc_context ERROR: %s in open_media \n", av_err2str(AVERROR(ENOMEM)));
        return AVERROR(ENOMEM);
    }
    //fprintf(stdout, "input file: %s \n", filename);
    if((ret = avformat_open_input(fmt_ctx, filename, NULL, NULL)) < 0){
        av_log(NULL, AV_LOG_ERROR, "avformat_open_input ERROR: %s in open_media \n", av_err2str(ret));
        return ret;
    }
    if((ret = avformat_find_stream_info(*fmt_ctx, NULL)) < 0){
        av_log(NULL, AV_LOG_ERROR, "avformat_find_stream_info ERROR: %s in open_media \n", av_err2str(ret));
        return ret;
    }
    return 0;
}

//准备每个流的AVCodecContex,AVCodec寻找信息
static int prepare_stream_info(AVCodecContext** p_codecctx, AVCodec** p_codec, AVStream* p_stm){
    int ret = 0;
    if((*p_codec = avcodec_find_decoder(p_stm->codecpar->codec_id)) == NULL){
        av_log(NULL, AV_LOG_ERROR, "avcodec_find_decoder ERROR in prepare_stream_info \n");
        return AVERROR(ENOMEM);
    }
    
    if((*p_codecctx = avcodec_alloc_context3(*p_codec)) == NULL){
        av_log(NULL, AV_LOG_ERROR, "avcodec_alloc_context3 ERROR in prepare_stream_info \n");
        return AVERROR(ENOMEM);
    }

    if((ret = avcodec_parameters_to_context(*p_codecctx, p_stm->codecpar)) < 0){
        av_log(NULL, AV_LOG_ERROR, "avcodec_parameters_to_context ERROR in prepare_stream_info \n");
        return ret;
    }

    if((ret = avcodec_open2(*p_codecctx, *p_codec, NULL)) < 0){
        av_log(NULL, AV_LOG_ERROR, "avcodec_open2 ERROR in prepare_stream_info \n");
        return ret;
    }
    return 0;
}

//输入文件和输出文件使用相同编解码器，直接复制信息
static int copy_stream(AVFormatContext *out_fmt_ctx, AVStream **out_stm, AVCodecParameters* in_codec_par){
    *out_stm = avformat_new_stream(out_fmt_ctx, NULL);
    if(*out_stm == NULL){
        av_log(NULL, AV_LOG_ERROR, "avformat_new_stream ERROR: in copy_stream \n");
        return AVERROR(EINVAL);
    }
    int ret = 0;
    if((ret = avcodec_parameters_copy((*out_stm)->codecpar, in_codec_par)) < 0){
        av_log(NULL, AV_LOG_ERROR, "avcodec_parameters_copy ERROR: in copy_stream \n");
        return AVERROR(EINVAL);
    }
    return 0;
}

// 准备编码的视频流信息
static int prepare_video_stream(MediaContext* out_mctx, AVCodecContext **in_codec_ctx, AVRational input_framerate, MediaDict *p_md){
    int ret = 0;
    out_mctx->v_stm = avformat_new_stream(out_mctx->fmt_ctx, NULL);
    if(out_mctx->v_stm == NULL){
        av_log(NULL, AV_LOG_ERROR, "avformat_new_stream ERROR in prepare_video_stream \n");
        return AVERROR(EINVAL);
    }

    out_mctx->v_codec = avcodec_find_encoder_by_name(p_md->video_codec);
    if(out_mctx->v_codec == NULL){
        av_log(NULL, AV_LOG_ERROR, "avcodec_find_encoder_by_name ERROR in prepare_video_stream \n");
        return AVERROR(EINVAL);
    }

    out_mctx->v_codec_ctx = avcodec_alloc_context3(out_mctx->v_codec);
    if(out_mctx->v_codec_ctx == NULL){
        av_log(NULL, AV_LOG_ERROR, "avcodec_alloc_context3 ERROR in prepare_video_stream \n");
        return AVERROR(EINVAL);
    }
    
    av_opt_set(out_mctx->v_codec_ctx->priv_data, "preset", "fast", 0);
    if(p_md->codec_opt_key && p_md->codec_opt_value){
        av_opt_set(out_mctx->v_codec_ctx->priv_data, p_md->codec_opt_key, p_md->codec_opt_value, 0);
    }

    out_mctx->v_codec_ctx->height = (*in_codec_ctx)->height;
    out_mctx->v_codec_ctx->width = (*in_codec_ctx)->width;
    out_mctx->v_codec_ctx->sample_aspect_ratio = (*in_codec_ctx)->sample_aspect_ratio;
    if(out_mctx->v_codec_ctx->pix_fmt){
        out_mctx->v_codec_ctx->pix_fmt = out_mctx->v_codec->pix_fmts[0];
    }else{
        out_mctx->v_codec_ctx->pix_fmt = (*in_codec_ctx)->pix_fmt;
    }

    out_mctx->v_codec_ctx->bit_rate = 2 * 1000 * 1000;
    out_mctx->v_codec_ctx->rc_buffer_size = 4 * 1024 * 1024;
    out_mctx->v_codec_ctx->rc_min_rate = 1 * 1000 * 1000;
    out_mctx->v_codec_ctx->rc_max_rate = 3 * 1000 * 1000;
    out_mctx->v_codec_ctx->time_base = av_inv_q(input_framerate);
    out_mctx->v_stm->time_base = out_mctx->v_codec_ctx->time_base;

    if((ret = avcodec_open2(out_mctx->v_codec_ctx, out_mctx->v_codec, NULL)) < 0){
        av_log(NULL, AV_LOG_ERROR, "avcodec_open2 ERROR in prepare_video_stream %s \n", av_err2str(ret));
        return AVERROR(ret);
    }
    
    if((ret = avcodec_parameters_from_context(out_mctx->v_stm->codecpar, out_mctx->v_codec_ctx)) < 0){
        av_log(NULL, AV_LOG_ERROR, "avcodec_parameters_from_context ERROR in prepare_video_stream \n");
        return AVERROR(ret);
    }

    return 0;
}

//
static int prepare_audio_stream(MediaContext* out_mctx, AVCodecContext **in_codec_ctx, int sample_rate, MediaDict *p_md){
    int ret = 0;
    out_mctx->a_stm = avformat_new_stream(out_mctx->fmt_ctx, NULL);
    if(out_mctx->a_stm == NULL){
        av_log(NULL, AV_LOG_ERROR, "avformat_new_stream ERROR in prepare_audio_stream \n");
        return AVERROR(EINVAL);
    }

    out_mctx->a_codec = avcodec_find_encoder_by_name(p_md->audio_codec);
    if(out_mctx->a_codec == NULL){
        av_log(NULL, AV_LOG_ERROR, "avcodec_find_encoder_by_name ERROR in prepare_audio_stream \n");
        return AVERROR(EINVAL);
    }

    out_mctx->a_codec_ctx = avcodec_alloc_context3(out_mctx->a_codec);
    if(out_mctx->a_codec_ctx == NULL){
        av_log(NULL, AV_LOG_ERROR, "avcodec_alloc_context3 ERROR in prepare_audio_stream \n");
        return AVERROR(EINVAL);
    }

    out_mctx->a_codec_ctx->channel_layout = av_get_default_channel_layout(2);
    out_mctx->a_codec_ctx->channels = 2;
    out_mctx->a_codec_ctx->sample_rate = sample_rate;
    out_mctx->a_codec_ctx->sample_fmt = out_mctx->a_codec->sample_fmts[0];
    out_mctx->a_codec_ctx->bit_rate = 196000;
    AVRational tb = {1, sample_rate};
    out_mctx->a_codec_ctx->time_base = tb;
    out_mctx->a_codec_ctx->strict_std_compliance = FF_COMPLIANCE_EXPERIMENTAL;
    out_mctx->a_stm->time_base = out_mctx->a_codec_ctx->time_base;
    
    if((ret = avcodec_open2(out_mctx->a_codec_ctx, out_mctx->a_codec, NULL)) < 0){
        av_log(NULL, AV_LOG_ERROR, "avcodec_open2 ERROR in prepare_audio_stream \n");
        return AVERROR(ret);
    }

    if((ret = avcodec_parameters_from_context(out_mctx->a_stm->codecpar, out_mctx->a_codec_ctx)) < 0){
        av_log(NULL, AV_LOG_ERROR, "avcodec_parameters_from_context ERROR in prepare_audio_stream \n");
        return AVERROR(ret);
    }

    return 0;
}

//准备解码器信息
static int prepare_decoder(MediaContext* m_ctx){
    int ret = 0;
    for(int i=0; i < m_ctx->fmt_ctx->nb_streams; i++){
        if(m_ctx->fmt_ctx->streams[i]->codecpar->codec_type == AVMEDIA_TYPE_VIDEO){
            m_ctx->v_stm = m_ctx->fmt_ctx->streams[i];
            m_ctx->v_idx = i;
            if((ret = prepare_stream_info(&m_ctx->v_codec_ctx, &m_ctx->v_codec, m_ctx->v_stm)) < 0){
                av_log(NULL, AV_LOG_ERROR, "prepare_stream_info ERROR in prepare_decoder \n");
                return ret;
            }
        }else if(m_ctx->fmt_ctx->streams[i]->codecpar->codec_type == AVMEDIA_TYPE_AUDIO){
            m_ctx->a_stm = m_ctx->fmt_ctx->streams[i];
            m_ctx->a_idx = i;
            if((ret = prepare_stream_info(&m_ctx->a_codec_ctx, &m_ctx->a_codec, m_ctx->a_stm)) < 0){
                av_log(NULL, AV_LOG_ERROR, "prepare_stream_info ERROR in prepare_decoder \n");
                return ret;
            }
        }else{
            av_log(NULL, AV_LOG_INFO, "No such media type! \n");
        }
    }
    return 0;
}

static int remux(AVFormatContext **fmt_ctx, AVPacket **p_packet, AVRational src_tb, AVRational des_tb){
    av_packet_rescale_ts(*p_packet, src_tb, des_tb);
    int ret = 0;
    if((ret = av_interleaved_write_frame(*fmt_ctx, *p_packet)) < 0){
        av_log(NULL, AV_LOG_ERROR, "av_interleaved_write_frame ERROR in remux \n");
        return AVERROR(ret);
    }
    return 0;

}

static int encoding_video(MediaContext *i_mctx, MediaContext *o_mctx, AVFrame *in_frame){

    AVPacket *out_packet = av_packet_alloc();
    if(out_packet == NULL){
        av_log(NULL, AV_LOG_ERROR, "av_packet_alloc ERROR in encoding_video \n");
        return AVERROR(ENOMEM);
    }
    int ret = 0;
    if((ret = avcodec_send_frame(o_mctx->v_codec_ctx, in_frame)) < 0){
        av_log(NULL, AV_LOG_ERROR, "avcodec_send_frame ERROR in encoding_video \n");
        return ret;
    }
    while(ret >= 0){
        ret = avcodec_receive_packet(o_mctx->v_codec_ctx, out_packet);
        if(ret == AVERROR(EAGAIN) || ret == AVERROR_EOF){
            break;
        }else if(ret < 0){
            av_log(NULL, AV_LOG_ERROR, "avcodec_receive_packet ERROR in encoding_video \n");
            return ret;
        }

        out_packet->stream_index = o_mctx->v_idx;
        out_packet->duration = i_mctx->v_stm->avg_frame_rate.num * i_mctx->v_stm->time_base.den / i_mctx->v_stm->avg_frame_rate.den / i_mctx->v_stm->time_base.num;
        av_packet_rescale_ts(out_packet, i_mctx->v_stm->time_base, o_mctx->v_stm->time_base);
        ret = av_interleaved_write_frame(o_mctx->fmt_ctx, out_packet);
        if(ret < 0){
            av_log(NULL, AV_LOG_ERROR, "av_interleaved_write_frame ERROR in encoding_video \n");
            return ret;
        }
    }
    av_packet_unref(out_packet);
    av_packet_free(&out_packet);
    return 0;
}

static int encoding_audio(MediaContext *i_mctx, MediaContext *o_mctx, AVFrame *in_frame){
    AVPacket *out_packet = av_packet_alloc();
    if(out_packet == NULL){
        av_log(NULL, AV_LOG_ERROR, "av_packet_alloc ERROR in encoding_audio \n");
        return AVERROR(ENOMEM);
    }
    int ret = 0;
    if((ret = avcodec_send_frame(o_mctx->a_codec_ctx, in_frame)) < 0){
        av_log(NULL, AV_LOG_ERROR, "avcodec_send_frame ERROR in encoding_audio \n");
        return ret;
    }
    while(ret >= 0){
        ret = avcodec_receive_packet(o_mctx->a_codec_ctx, out_packet);
        if(ret == AVERROR(EAGAIN) || ret == AVERROR_EOF){
            break;
        }else if(ret < 0){
            av_log(NULL, AV_LOG_ERROR, "avcodec_receive_packet ERROR %s in encoding_audio \n", av_err2str(ret));
            return ret;
        }

        out_packet->stream_index = o_mctx->a_idx;
        
        av_packet_rescale_ts(out_packet, i_mctx->a_stm->time_base, o_mctx->a_stm->time_base);
        ret = av_interleaved_write_frame(o_mctx->fmt_ctx, out_packet);
        if(ret < 0){
            av_log(NULL, AV_LOG_ERROR, "av_interleaved_write_frame ERROR %s in encoding_audio \n", av_err2str(ret));
            return ret;
        }
    }
    av_packet_unref(out_packet);
    av_packet_free(&out_packet);
    return 0;
}

static int transcoding_video(MediaContext *i_mctx, MediaContext *o_mctx, AVPacket *p_packet, AVFrame *p_frame){
    int ret = 0;
    if((ret = avcodec_send_packet(i_mctx->v_codec_ctx, p_packet)) < 0){
        av_log(NULL, AV_LOG_ERROR, "avcodec_send_packet ERROR in transcoding_video \n");
        return AVERROR(ret);
    }
    while(ret >= 0){
        ret = avcodec_receive_frame(i_mctx->v_codec_ctx, p_frame);
        if(ret == AVERROR(EAGAIN) || ret == AVERROR_EOF){
            break;
        }else if(ret < 0){
            av_log(NULL, AV_LOG_ERROR, "avcodec_receive_frame ERROR in transcoding_video \n");
            return ret;
        }
        if((ret = encoding_video(i_mctx, o_mctx, p_frame)) < 0){
            return ret;
        }
        av_frame_unref(p_frame);
    }
    return 0;
}

static int transcoding_audio(MediaContext *i_mctx, MediaContext *o_mctx, AVPacket *p_packet, AVFrame *p_frame){
    int ret = 0;
    if((ret = avcodec_send_packet(i_mctx->a_codec_ctx, p_packet)) < 0){
        av_log(NULL, AV_LOG_ERROR, "avcodec_send_packet ERROR in transcoding_audio \n");
        return AVERROR(ret);
    }
    while(ret >= 0){
        ret = avcodec_receive_frame(i_mctx->a_codec_ctx, p_frame);
        if(ret == AVERROR(EAGAIN) || ret == AVERROR_EOF){
            break;
        }else if(ret < 0){
            av_log(NULL, AV_LOG_ERROR, "avcodec_receive_frame ERROR %s in transcoding_audio \n", av_err2str(ret));
            return ret;
        }
        if((ret = encoding_audio(i_mctx, o_mctx, p_frame)) < 0){
            return ret;
        }
        av_frame_unref(p_frame);
    }
    return 0;
}

int main(int argc, char* argv[]){
    
    if(argc < 3){
        av_log(NULL, AV_LOG_ERROR, "You must specify the input file and the output file!\n");
        return -1;
    }
    

    fprintf(stdout, "Input media file : %s, Output media file : %s \n", argv[1], argv[2]);
    
    MediaDict mdict;
    mdict.copy_audio = 1;
    mdict.copy_video = 0;
    mdict.video_codec = "libx265"; 
    mdict.muxer_opt_key = NULL;
    mdict.muxer_opt_value = NULL;

    MediaContext *input_mctx = NULL;
    MediaContext *output_mctx = NULL;

    input_mctx = av_malloc_array(1, sizeof(*input_mctx));
    output_mctx = av_malloc_array(1, sizeof(*output_mctx));
    input_mctx->name = argv[1];
    output_mctx->name = argv[2];

    int ret = 0;
    if((ret = open_media(&(input_mctx->fmt_ctx), input_mctx->name)) < 0){
        av_log(NULL, AV_LOG_ERROR, "open_media ERROR: %s in main\n", av_err2str(ret));
        return ret;
    }

    if((ret = prepare_decoder(input_mctx)) < 0){
        av_log(NULL, AV_LOG_ERROR, "prepare_decoder ERROR: %s in main\n", av_err2str(ret));
        return ret;
    }
    
    avformat_alloc_output_context2(&output_mctx->fmt_ctx, NULL, NULL, output_mctx->name);
    if(output_mctx->fmt_ctx == NULL){
        av_log(NULL, AV_LOG_ERROR, "avformat_alloc_output_context2 ERROR: %s in main\n", av_err2str(AVERROR(ENOMEM)));
        return AVERROR(ENOMEM);
    }

    if(mdict.copy_video){
        copy_stream(output_mctx->fmt_ctx, &(output_mctx->v_stm), input_mctx->v_stm->codecpar);
    }else{
        AVRational input_framerate = av_guess_frame_rate(input_mctx->fmt_ctx, input_mctx->v_stm, NULL);
        prepare_video_stream(output_mctx, &(input_mctx->v_codec_ctx), input_framerate, &mdict);
    }
    
    if(mdict.copy_audio){
        copy_stream(output_mctx->fmt_ctx, &(output_mctx->a_stm), input_mctx->a_stm->codecpar);
    }else{
        prepare_audio_stream(output_mctx, &(input_mctx->v_codec_ctx), input_mctx->a_codec_ctx->sample_rate, &mdict);
    }

    if(output_mctx->fmt_ctx->oformat->flags & AVFMT_GLOBALHEADER){
        output_mctx->fmt_ctx->flags |= AV_CODEC_FLAG_GLOBAL_HEADER;
    }

    if(!(output_mctx->fmt_ctx->oformat->flags & AVFMT_NOFILE)){
        if((ret = avio_open2(&(output_mctx->fmt_ctx->pb), output_mctx->name, AVIO_FLAG_WRITE, NULL, NULL)) < 0){
            av_log(NULL, AV_LOG_ERROR, "avio_open2 ERROR %s in main\n", av_err2str(ret));
            return AVERROR(EINVAL);
        }
    }

    AVDictionary* muxer_dict = NULL;
    if(mdict.muxer_opt_key != NULL && mdict.muxer_opt_value != NULL){
        fprintf(stdout, "key: %s, value: %s", mdict.muxer_opt_key, mdict.muxer_opt_value);
        if ((ret = av_dict_set(&muxer_dict, mdict.muxer_opt_key, mdict.muxer_opt_value, 0)) < 0){
            av_log(NULL, AV_LOG_ERROR, "av_dict_set ERROR %s in main \n", av_err2str(ret));
            return AVERROR(ret);
        }
    }

    if((ret = avformat_write_header(output_mctx->fmt_ctx, &muxer_dict))){
        av_log(NULL, AV_LOG_ERROR, "avformat_write_header ERROR %s in main \n", av_err2str(ret));
        return AVERROR(ret);
    }
    
    AVPacket* p_packet = av_packet_alloc();
    if(p_packet == NULL){
        av_log(NULL, AV_LOG_ERROR, "av_packet_alloc ERROR in main \n");
        return AVERROR(ENOMEM);
    }

    AVFrame* p_frame = av_frame_alloc();
    if(p_frame == NULL){
        av_log(NULL, AV_LOG_ERROR, "av_frame_alloc ERROR in main \n");
        return AVERROR(ENOMEM);
    }

    while((ret = av_read_frame(input_mctx->fmt_ctx, p_packet)) == 0){
        if(input_mctx->fmt_ctx->streams[p_packet->stream_index]->codecpar->codec_type == AVMEDIA_TYPE_VIDEO){
            if (mdict.copy_video){
                remux(&output_mctx->fmt_ctx, &p_packet, input_mctx->v_stm->time_base, output_mctx->v_stm->time_base);
            } else{
                if((ret = transcoding_video(input_mctx, output_mctx, p_packet, p_frame)) < 0){
                    av_log(NULL, AV_LOG_ERROR, "encoding_video ERROR %s in main\n", av_err2str(ret));
                    return ret;
                } 
            }
        } else if(input_mctx->fmt_ctx->streams[p_packet->stream_index]->codecpar->codec_type == AVMEDIA_TYPE_AUDIO){
            if (mdict.copy_audio){
                remux(&output_mctx->fmt_ctx, &p_packet, input_mctx->a_stm->time_base, output_mctx->a_stm->time_base);
            } else{
                if((ret = transcoding_audio(input_mctx, output_mctx, p_packet, p_frame)) <0){
                    av_log(NULL, AV_LOG_ERROR, "encoding_audio ERROR %s in main\n", av_err2str(ret));
                    return ret;
                }
            }
        } else{
            av_log(NULL, AV_LOG_INFO, "NO Such Type Media\n");
        }
        
    }

    encoding_video(input_mctx, output_mctx, NULL);

    av_write_trailer(output_mctx->fmt_ctx);
    
    if(muxer_dict != NULL){
        av_dict_free(&muxer_dict);
        muxer_dict = NULL;
    }
    if(p_frame != NULL){
        av_frame_free(&p_frame);
        p_frame = NULL;
    }
    if(p_packet != NULL){
        av_packet_free(&p_packet);
        p_packet = NULL;
    }
    avformat_close_input(&output_mctx->fmt_ctx);
    avformat_free_context(input_mctx->fmt_ctx);
    input_mctx->fmt_ctx = NULL;
    avformat_free_context(output_mctx->fmt_ctx);
    output_mctx->fmt_ctx = NULL;

    if(output_mctx->v_codec_ctx != NULL){
        avcodec_free_context(&(output_mctx->v_codec_ctx));
    }
    if(output_mctx->a_codec_ctx != NULL){
        avcodec_free_context(&(output_mctx->a_codec_ctx));
    }
    
    av_freep(input_mctx);
    input_mctx = NULL;
    av_freep(output_mctx);
    output_mctx = NULL;

    return 0;
}